#include <map>
#include <stdlib.h>
#include <errno.h>
#include <string.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <sys/poll.h>
#include <time.h>
#include <getopt.h>
#include <string>
#include <iostream>
#include <stdio.h>
#include <mysql/mysql.h>

const char ConfigFile[] = "config";
static const int timeout = 1;
static const int OPEN_MAX = 10;

using namespace std;

map<int, struct sockaddr_in> id2Addr;

void buildAddress(struct sockaddr_in* addr, string address, int port)
{
  bzero(addr, sizeof(struct sockaddr));
  addr->sin_family = AF_INET;
  addr->sin_port = htons(port);
  if(address == "localhost")
    addr->sin_addr.s_addr = htonl(INADDR_ANY);
  else
    inet_pton(AF_INET, address.c_str(), &addr->sin_addr);
}

int readn(int fd, char* buff, int len)
{
  int nleft = len;
  int nread;
  char *ptr = buff;
  while(nleft>0)
  {
    nread = read(fd, buff, nleft);
	if(nread<0)
	{
	  if(errno == EINTR)
	    nread = 0;
      else
	    return -1;
	}
	else if(nread == 0)
	  break;
    nleft -= nread;
    ptr += nread;
  }

  return len-nleft;
}

int UDPServer(string address, int port)
{
  int res;

  /* Create UDP server socket */
  int usock = socket(AF_INET, SOCK_DGRAM, 0);
  if(-1 == usock)
  {
    cerr<<"Create socket for receive heartbeat error: "<<strerror(errno)<<endl;
    exit(1);
  }
  int reuse = 1;
  setsockopt(usock, SOL_SOCKET, SO_REUSEADDR, (const char*)&reuse, sizeof(reuse));

  /* Build local address */
  struct sockaddr_in addr;
  buildAddress(&addr, address, port);

  res = bind(usock, (struct sockaddr*)&addr, sizeof(struct sockaddr));
  cout<<"UDP server bind on: "<<ntohs(addr.sin_port)<<endl;
  if(-1 == res)
  {
    cerr<<"UDP server bind error: "<<strerror(errno)<<endl;
    exit(1);
  }

  return usock;
}

int TCPServer(string address, int port)
{
  int res;
  int reuse = 1;

  /* Create TCP server socket */
  int tsock = socket(AF_INET, SOCK_STREAM, 0);
  if(-1 == tsock)
  {
    cerr<<"Create socket for receive heartbeat error: "<<strerror(errno)<<endl;
    exit(1);
  }

  setsockopt(tsock, SOL_SOCKET, SO_REUSEADDR, (const char*)&reuse, sizeof(reuse));

  /* Build local address */
  struct sockaddr_in addr;
  buildAddress(&addr, address, port);

  res = bind(tsock, (struct sockaddr*)&addr, sizeof(struct sockaddr));
  cout<<"TCP server bind on: "<<ntohs(addr.sin_port)<<endl;
  if(-1 == res)
  {
    cerr<<"UDP server bind error: "<<strerror(errno)<<endl;
    exit(1);
  }
  
  listen(tsock, 5);
  
  return tsock;
}

int sendUDP(int sock, string address, int port, char* buff, int size)
{
  /* Build address */
  struct sockaddr_in addr;
  buildAddress(&addr, address, port);
  return sendto(sock, buff, size, 0, (struct sockaddr*)&addr, sizeof(struct sockaddr));
}

int main(int argc, char** argv)
{
  int c;
  int helpFlag = 0,
      errFlag = 0,
      daemonFlag = 0;
  
  int uport = 0, 
      tport = 0;

  struct option longopts[] = 
  {
    {"daemon", 0, 0, 'd'},
    {"uport", required_argument, 0, 'u'},
    {"tport", required_argument, 0, 't'},
    {0, 0, 0, 0}
  };
  
  while((c = getopt_long(argc, argv, "t:u:d", longopts, NULL)) != EOF)
  {
    switch(c)
    {
      case 'd':
        daemonFlag = 1;
        break;
      case 'u':
        uport = atoi(optarg);
        break;
      case 't':
        tport = atoi(optarg);
        break;
      case '?':
        errFlag++;
        break;
      default:
        break;
    }
  }
  
  if(daemonFlag)
    daemon(0, 0);
  
  if(uport == 0 || tport == 0)
    errFlag++;
 
  if(errFlag)
  {
    cerr<<"Usage:  heartbeatServer [OPTION]"<<endl<<endl
        <<"\t-t, --tport  port for receive data from control server"<<endl
        <<"\t-u, --uport  port for receive hearbeat from client"<<endl
        <<"\t-d, --daemon set program to daemon"<<endl<<endl;
    exit(0);
  }

  int res;

  struct sockaddr_in cliaddr;
  socklen_t clilen = sizeof(struct sockaddr_in);
  int connfd;
  
  int usock = UDPServer(string("localhost"), uport);
  int tsock = TCPServer(string("localhost"), tport);

  struct pollfd fds[10];
	for(int i=0;i<10;++i)
		fds[i].fd = -1;
  int maxi = 1; /* max index into fds[] array */
  fds[0].fd = tsock;
  fds[0].events = POLLIN;

  fds[1].fd = usock;
  fds[1].events = POLLIN;

  char buff[1024];
  time_t now;
  struct tm *timenow;

	MYSQL mysql;
	char sql[1000];
	mysql_init(&mysql);
	mysql_real_connect(&mysql, "localhost", "root", "yolein", "heartbeat", 3306, NULL, 0);

  while(true)
  {
    time(&now);
    timenow = localtime(&now);

//		cout<<"maxi: "<<maxi<<endl;
    res = poll(fds, maxi + 1, 1000 * timeout);
//		cout<<res<<endl;

		bzero(buff, sizeof(buff));

		if(fds[0].revents & POLLIN) /* TCP server */
		{
      connfd = accept(tsock, (struct sockaddr*)&cliaddr, &clilen);

	  	cout<<"new connect: "<<connfd<<endl;

		  int i;
		  for(i = 2; i < OPEN_MAX; i++)
	  	{
        if(fds[i].fd < 0)
				{
          fds[i].fd = connfd;
  				fds[i].events = POLLIN;
			    break;
				}
	  	}

		  if(OPEN_MAX == i)
		  {
        cerr<<"Too many clients"<<endl;
	  	  exit(1);
		  }
		  if(i > maxi)
	  	  maxi = i;
		  if(--res <= 0)
		    continue;
		}

    if(fds[1].revents & POLLIN) //UDP server
		{
		  cout<<"------------------- UDP start -------------------------"<<endl;
      cout<<"Time: "<<asctime(timenow)<<endl;
      res = recvfrom(usock, buff, sizeof(buff), 0, (struct sockaddr*)&cliaddr, &clilen);
	  	if(-1 == res)
		  {
		    cout<<strerror(errno)<<endl;
	  	  cout<<"*"<<endl;
		  }

      //print package
		  for(int j=0;j<res;j++)
	  	  printf("%02hhx ", buff[j]);
		  cout<<endl;
	  
		  int id;
	  	if(buff[0] == '\x01') //command package
		  {
        id = int(buff[1]); 
        cout<<"Receive command "<<endl
						<<"ID: "<<int(buff[1])<<endl
				    <<"Type: "<<int(buff[2])<<endl
				    <<"Send to "<<inet_ntoa(id2Addr[id].sin_addr)<<endl
		    		<<"Port: "<<ntohs(id2Addr[id].sin_port)<<endl;
				sendto(fds[1].fd, buff + 2, res - 2, 0, (struct sockaddr*)&id2Addr[id], sizeof(id2Addr[id]));
	  	}
		  else if(buff[0] == '\x02') //heartbeat package
		  {
        sscanf(buff + 1, "%d", &id);
				sprintf(sql, "insert into hb values (%d, '%d-%d-%d %d:%d:%d');", id, timenow->tm_year + 1900, timenow->tm_mon + 1, timenow->tm_mday, timenow->tm_hour, timenow->tm_min, timenow->tm_sec);
				mysql_query(&mysql, sql);
	  	  id2Addr[id] = cliaddr;
	    	cout<<"Recieve heartbeat  "<<id<<endl
	      	  <<"Address: "<<inet_ntoa(id2Addr[id].sin_addr)<<endl
	        	<<"Port : "<<ntohs(id2Addr[id].sin_port)<<endl
		        <<"Heartbeat ID: "<<id<<endl;
		  }
	  	else
	    	cout<<"Bad package"<<endl;
		  cout<<"------------------- UDP end -------------------------"<<endl;
		}

		for(int i = 2; i <= maxi; i++)
		{
  		if(fds[i].fd < 0)
	    	continue;
		  if(fds[i].revents & POLLIN)
		  {
	  	  cout<<"------------------- TCP connect start -------------------------"<<endl;
      	cout<<"Time: "<<asctime(timenow)<<endl;
		    char len;
	      res = readn(fds[i].fd, &len, sizeof(len));
		    cout<<"res: "<<res<<endl;
	  	  if(0 == res)
	    	{
        	close(fds[i].fd);
		      fds[i].fd = -1;
				  cout<<"Disconnected"<<endl;
	  	    cout<<"------------------- TCP connect end -------------------------"<<endl;
		  		continue;
				}

				res = readn(fds[i].fd, buff + 1, len);
				cout<<"Read len: "<<res<<endl;
	
				for(int j=0;j< res;j++)
				  printf("\\x%02hhx ", buff[j + 1]);
				cout<<endl;
	
				if(res == 0)
				{
      	  close(fds[i].fd);
		  		fds[i].fd = -1;
				}
				else
				{
				  buff[0] = '\x01'; //set flag to cmd
			  	/* Send command to UDP server then forward to camera control*/
	        sendUDP(usock, string("localhost"), uport, buff, res + 1);
				  cout<<"Forward"<<endl;
				} 
		    cout<<"------------------- TCP connect end -------------------------"<<endl;
	  	}
		}
	}
}
